package main

import (
	"context"
	"flag"
	"fmt"

	"github.com/zeromicro/go-zero/core/conf"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/mcp"
)

var configFile = flag.String("f", "etc/config.yaml", "配置文件")

func main() {
	flag.Parse()

	var c mcp.McpConf
	conf.MustLoad(*configFile, &c)

	// 创建并启动 MCP 服务器
	server := mcp.NewMcpServer(c)

	// 注册你的 Prompts 处理器
	registerMyPrompt(server)

	// 启动服务器
	logx.Info("正在启动带有 SSE 传输的 MCP 服务器")
	server.Start()
}

// registerMyPrompt 注册动态 Prompts 处理器
func registerMyPrompt(server mcp.McpServer) {
	server.RegisterPrompt(mcp.Prompt{
		Name:        "get_custom_prompt",
		Description: "基于输入参数生成自定义提示词",
		Arguments: []mcp.PromptArgument{
			{
				Name:        "user_query",
				Description: "用户的原始查询",
				Required:    true,
			},
			{
				Name:        "context_id",
				Description: "可选的上下文标识符",
				Required:    false,
			},
		},
		Handler: func(ctx context.Context, args map[string]string) ([]mcp.PromptMessage, error) {
			// 解析参数
			var req struct {
				UserQuery string `json:"user_query"`
				ContextID string `json:"context_id,optional"`
			}
			if err := mcp.ParseArguments(args, &req); err != nil {
				return nil, fmt.Errorf("参数解析失败: %w", err)
			}

			// 基于输入参数生成动态提示词
			return []mcp.PromptMessage{
				{
					Role: mcp.RoleUser,
					Content: mcp.TextContent{
						Text: fmt.Sprintf(`你是一个有帮助的助手。请回答以下问题: %s %s`, req.UserQuery, getContextInstructions(req.ContextID)),
					},
				},
			}, nil
		},
	})
}

func getContextInstructions(contextID string) string {
	if contextID == "" {
		return ""
	}

	// 在实际应用中，你可能会基于 contextID 从数据库或外部服务
	// 获取特定上下文的指令
	return fmt.Sprintf("\n回答时，请考虑 ID 为 %s 的上下文", contextID)
}
